;; Non-robust row reduced echelon form functions
;; Copyright (C) 2025  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.sole
  (:require
   [mazama.linalg.rref :as rref]))

(defn- rows-as-data
  [rows]
  (apply concat rows))

(defn- column [j rows]
  (map #(nth % j) rows))

(defn- nullvector
  "Compute an element of the nullspace."
  [j acols pci pivot-rows]
  (loop [k 0
         c (column j pivot-rows)
         pci (seq pci)
         result []]
    (cond
      (>= k acols) result
      (== k j) (recur (inc k) c pci (conj result -1))
      (and pci (== k (first pci))) (recur (inc k) (rest c) (next pci) (conj result (first c)))
      :else (recur (inc k) c pci (conj result 0)))))

(defn- de-augment-b
  "Return b from the augmented matrix A|b as a sequence of rows."
  [acols rows]
  (map #(drop acols %) rows))

(defn- general-solution
  "Solve Ax=b, assuming A|b is in rref."
  [ncols acols pivot-column-indexes pivot-rows]
  (let [acols (or acols ncols)
        b-zeros (repeat (- ncols acols) 0)]
    (loop [j 0
           pci-past []
           pci-future (seq pivot-column-indexes)
           b (de-augment-b acols pivot-rows)
           particular-solution []
           nullspace []]
      (cond
        (>= j acols) {:ncols (- ncols acols)
                      :data (rows-as-data particular-solution)
                      :nullspace nullspace}
        (and pci-future
             (== j (first pci-future))) (recur (inc j)
                                               (conj pci-past (first pci-future))
                                               (next pci-future)
                                               (rest b)
                                               (conj particular-solution (first b))
                                               nullspace)
        :else (recur (inc j)
                     pci-past
                     pci-future
                     b
                     (conj particular-solution b-zeros)
                     (conj nullspace (nullvector j acols pci-past pivot-rows)))))))

(defn- one-solution
  "Solve Ax=b, assuming A|b is in rref and A only contains pivot columns."
  [acols pivot-rows]
  {:ncols (- (count (first pivot-rows)) acols)
   :data (rows-as-data (de-augment-b acols pivot-rows))})

(defn- inconsistent?
  "If any zero-row contains a non-zero."
  [zero-rows]
  (let [f (fn [_ row]
            (if (every? zero? row)
              false
              (reduced true)))]
    (reduce f false zero-rows)))

(defn- augment
  "Augment two matrices `A` and `b` to become A|b."
  [A b]
  (let [{acols :ncols adata :data} A
        {bcols :ncols bdata :data} b]
    (loop [adata adata
           bdata bdata
           data []]
      (if (or (empty? adata) (empty? bdata))
        {:ncols (+ acols bcols) :data data :acols acols}
        (recur (drop acols adata)
               (drop bcols bdata)
               (into data (concat (take acols adata) (take bcols bdata))))))))

(defn solve
  "Solve a system of linear equations, Ax=b.
  `A` is an mxn matrix.
  `b` is an mxp matrix, for some p > 0.
  If A and b do not have the same number of rows then the result is undefined,
  otherwise the result is:
  when there is no solution: nil;
  when there is a single solution: an nxp matrix;
  when there are infinitely many solutions: an nxp matrix with an additional
  property: :nullspace a-collection-of-nullspace-basis-vectors."
  [A b]
  (let [{:keys [ncols acols] :as m} (augment A b)
        {:keys [pivot-column-indexes pivot-rows zero-rows]} (rref/rref m)]
    (cond
      (inconsistent? zero-rows) nil
      (== (count pivot-column-indexes) acols) (one-solution acols pivot-rows)
      :else (general-solution ncols acols pivot-column-indexes pivot-rows))))

(defn nullspace
  "Return the nullspace of the matrix `m`, as a collection of nullspace basis vectors."
  [m]
  (let [{:keys [pivot-column-indexes pivot-rows]} (rref/rref m)]
    (:nullspace (general-solution (:ncols m) nil pivot-column-indexes pivot-rows))))
